use base64::prelude::*;
use std::time::{SystemTime, UNIX_EPOCH};
use sha1::{Digest, Sha1};

const DEFAULT_EXPIRE_DURATION: u64 = 30;
const DEFAULT_CODE_LENGTH: usize = 6;

pub fn get_time_count(expire_duration: Option<u64>) -> u64 {
    let expire_duration = expire_duration.unwrap_or(DEFAULT_EXPIRE_DURATION);
    let secs = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs();
    secs / expire_duration
}

pub fn generate_totp(secret: &str, time_count: u64, code_length: Option<usize>) -> String {
    let code_length = code_length.unwrap_or(DEFAULT_CODE_LENGTH);
    let decoded_secret = BASE64_STANDARD.decode(secret).unwrap();
    let mut hmac = Sha1::new();
    hmac.update(&decoded_secret);
    hmac.update(&time_count.to_be_bytes());
    let result = hmac.finalize();
    
    let offset = (result[19] & 0xf) as usize;
    let code = u32::from_be_bytes([
        result[offset] & 0x7f,
        result[offset + 1],
        result[offset + 2],
        result[offset + 3],
    ]);
    
    format!("{:0width$}", code % 10u32.pow(code_length as u32), width = code_length)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_get_time_count() {
        let time_count = get_time_count(None);
        assert!(time_count > 0);
        
        let time_count2 = get_time_count(None);
        assert!(time_count2.saturating_sub(time_count) <= 1 || time_count.saturating_sub(time_count2) <= 1);

        let custom_expire_duration = Some(60);
        let time_count3 = get_time_count(custom_expire_duration);
        assert!(time_count3 > 0);
    }

    #[test]
    fn test_generate_totp() {
        let secret = "JBSWY3DPEHPK3PXP"; // 示例密钥
        let time_count = 1234567890;
        let totp = generate_totp(secret, time_count, None);
        
        assert_eq!(totp.len(), DEFAULT_CODE_LENGTH);
        assert!(totp.chars().all(|c| c.is_digit(10)));
        
        let totp2 = generate_totp(secret, time_count, None);
        assert_eq!(totp, totp2);

        let totp3 = generate_totp(secret, time_count + 1, None);
        assert_ne!(totp, totp3);

        let custom_length = Some(8);
        let totp4 = generate_totp(secret, time_count, custom_length);
        assert_eq!(totp4.len(), 8);
    }

    #[test]
    fn test_totp_expiration() {
        let secret = "JBSWY3DPEHPK3PXP"; // 示例密钥
        let expire_duration = Some(30);
        
        let time_count = get_time_count(expire_duration);
        let totp1 = generate_totp(secret, time_count, None);
        
        std::thread::sleep(std::time::Duration::from_secs(10));
        
        let new_time_count = get_time_count(expire_duration);
        let totp2 = generate_totp(secret, new_time_count, None);
        
        assert_eq!(totp1, totp2);
        assert_eq!(time_count, new_time_count);
        
        std::thread::sleep(std::time::Duration::from_secs(21));
        
        let expired_time_count = get_time_count(expire_duration);
        let totp3 = generate_totp(secret, expired_time_count, None);
        
        assert_ne!(totp1, totp3);
        assert_ne!(time_count, expired_time_count);
    }
}